#ifndef _GMG_H_
#define _GMG_H_

#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>
#include <Eigen/Dense>
#include <ctime>
#include <cmath>
#include <string>
#include "Equation.h"
#include "Mesh.h"
#include "function2.h"
#define POW2(a) (1 << a)

using namespace Eigen;
using namespace std;
typedef Eigen::SparseMatrix<double> SpMat;
typedef Eigen::VectorXd RHS;
typedef Eigen::VectorXd Solution;
typedef std::string Method;

template <unsigned int DIM>
std::map<int, double> interpolateBoundaryValues(auto &_boundary_function, std::set<Dofs<DIM>> _boundary_dofs, std::vector<Point<DIM>> _vertex)
{
    std::map<int, double> boundary_values;
    for (auto &dofs : _boundary_dofs)
    {
        int index = dofs.getGlobalIndex();
        double p[DIM];
        for (int i = 0; i < DIM; i++)
            p[i] = dofs[i];
        double value;
        if (p[1] == _vertex[0][1])
            value = _boundary_function[0](p);
        else if (p[0] == _vertex[1][0])
            value = _boundary_function[1](p);
        else if (p[1] == _vertex[1][1])
            value = _boundary_function[2](p);
        else if (p[0] == _vertex[0][0])
            value = _boundary_function[3](p);
        else
            abort(3);
        boundary_values.insert({index, value});
    }
    return boundary_values;
}
/**
 * @brief 对刚度矩阵和右端项应用Dirichlet条件
 *
 * @param system_matrix 刚度矩阵
 * @param system_rhs 右端项
 */

void applyBoundaryValues(std::map<int, double> &boundary_values, SpMat &system_matrix, RHS &system_rhs)
{
    for (auto &it : boundary_values)
    {
        // 边界自由度的全局编号。
        int k = it.first;
        // 对应自由的边界函数值。
        double bnd_value = it.second;
        system_rhs[k] = bnd_value * system_matrix.coeffRef(k, k);
        for (SpMat::InnerIterator it(system_matrix, k); it; ++it)
        {
            int row = it.row();
            if (row == k)
                continue;
            system_rhs[row] -= system_matrix.coeffRef(k, row) * bnd_value;
            system_matrix.coeffRef(k, row) = 0.0;
            system_matrix.coeffRef(row, k) = 0.0;
        }
    }
}

class MultiGridMethod
{
protected:
    vector<SpMat> A;    // 所有A矩阵
    vector<Solution> u; // 所有解向量
    vector<RHS> f;      // 所有右端项

    int level;    // 当前层数
    int maxlevel; // 最大层数
    int num;      // 松弛次数
    Eigen::ConjugateGradient<SpMat> Solver_sparse;

    Mesh<2> *mesh;

public:
    MultiGridMethod(){};
    ~MultiGridMethod(){};
    void setMaxLevel(int _n) { maxlevel = _n; }
    void setNum(int _n) { num = _n; }
    void Iterator(SpMat &_A, Solution &_u, RHS &_f, const Method &_m, Mesh<2> *_mesh);
    void Smooth(SpMat &_A, Solution &_u, RHS &_f, const Method &_m, int _n, Mesh<2> *_mesh);
    virtual void V_Cycle(int _level) = 0;
    Solution getSolution() { return u[0]; };
    void DealWithBoundaryCondition(SpMat &A, VectorXd &Rhs);
    void GS(SpMat &_A, Solution &_u, RHS &_f, Mesh<2> *_mesh);
    VectorXd &Solve();
    int switcher(int i, int j, int nx, int ny);
};

VectorXd &MultiGridMethod::Solve()
{
    std::cout << "Solve..." << std::endl;
    int times = 0;
    VectorXd residual_vector = A[0] * Eigen::MatrixXd::Zero(f[0].size(), 1) - f[0];
    // VectorXd residual_vector = _A_origin * Eigen::MatrixXd::Zero(_Rhs[0].size(), 1) - _Rhs_origin;
    double residual = residual_vector.norm();
    double residual_max = residual_vector.lpNorm<Eigen::Infinity>();
    do
    {
        V_Cycle(1);
        f.resize(1);
        u.resize(1);
        SpMat _A_temp = A[0];
        VectorXd _Rhs_temp = f[0];
        level = 1;
        DealWithBoundaryCondition(_A_temp, _Rhs_temp);
        residual_vector = _A_temp * u[0] - _Rhs_temp;
        residual = residual_vector.norm();
        residual_max = residual_vector.lpNorm<Eigen::Infinity>();
        std::cout << " the " << times + 1 << "th l2 norm of residual vecror is :" << residual << " and the max norm case is :" << residual_max << std::endl;
    } while (times++ < 10 & residual >= 1e-6);
    return u[0];
}

void MultiGridMethod::GS(SpMat &_A, Solution &_u, RHS &_f, Mesh<2> *_mesh)
{
    int index = 0;
    vector<int> indexofBoundaryDofs;
    for (int i = 0; i < _mesh->getSegment()[0] + 1; i++)
        indexofBoundaryDofs.push_back(i);

    for (int j = 1; j < _mesh->getSegment()[1]; j++)
    {
        indexofBoundaryDofs.push_back(j * (_mesh->getSegment()[0] + 1));
        indexofBoundaryDofs.push_back((j + 1) * (_mesh->getSegment()[0] + 1) - 1);
    }

    for (int i = 0; i < _mesh->getSegment()[0] + 1; i++)
        indexofBoundaryDofs.push_back(i + (_mesh->getSegment()[0] + 1) * _mesh->getSegment()[1]);

    for (int i = 0; i < _u.size(); i++)
    {
        double f_i = _f[i];
        double a_ii = _A.coeffRef(i, i);
        if (i == indexofBoundaryDofs[index])
        {
            // if (_f.size() == (_mesh->getSegment()[0] + 1) * (_mesh->getSegment()[1] + 1))
            if (level == 1)
            {
                _u[i] = f_i / a_ii;
                index++;
                // for (int i = 0; i < _u.size(); i++)
                // cout << i << "  " << _u[i] << endl;
                // cout << "hhhh" << endl;
            }
        }
        else
        {
            for (SpMat::InnerIterator it(_A, i); it; ++it)
            {
                int row = it.row();
                if (row == i)
                    continue;
                else
                    f_i -= (it.value() * _u[row]); // 把u写进函数参数，方便计算后自动更新u带入下一步计算
            }
            _u[i] = f_i / a_ii;
        }
    }
}

void MultiGridMethod::Iterator(SpMat &_A, Solution &_u, RHS &_f, const Method &_m, Mesh<2> *_mesh)
{
    if (_m.compare("GS") == 0)
        GS(_A, _u, _f, _mesh);
}

void MultiGridMethod::Smooth(SpMat &_A, Solution &_u, RHS &_f, const Method &_m, int _n, Mesh<2> *_mesh)
{
    for (int i = 0; i < _n; i++)
        Iterator(_A, _u, _f, _m, _mesh);
}

void MultiGridMethod::DealWithBoundaryCondition(SpMat &_A, VectorXd &_Rhs)
{
    vector<Point<2>> p = {{-1.0, -1.0}, {1.0, -1.0}, {1.0, 1.0}, {-1.0, 1.0}};
    Domain<2> *domain = new Rectangle2D(p);
    vector<int> seg({mesh->getSegment()[0] / POW2(level - 1), mesh->getSegment()[1] / POW2(level - 1)});
    Q1Mesh mesh_t(domain, seg);

    Equation<2> *equ = new PossionEquation<2>();
    equ->SetRightHandsTermFunction(fe);
    equ->SetBoundaryConditionFunction({g1, g2, g3, g4});

    std::set<Dofs<2>> boundary_dofs;
    boundary_dofs = mesh_t.getBoundaryDofs();
    std::map<int, double> boundary_values;
    auto boundary_function = equ->BoundaryConditionFunction();
    std::vector<Point<2>> vertex = {domain->getVertex(0), domain->getVertex(2)};
    boundary_values = interpolateBoundaryValues(boundary_function, boundary_dofs, vertex);
    // applyBoundaryValues(boundary_values, _A, _Rhs);

    for (auto &it : boundary_values)
    {
        int k = it.first;
        double bnd_value = it.second;
        // if (level == 1)
        //     _Rhs[k] = bnd_value * _A.coeffRef(k, k);
        // level = 1已经在Femspace中的dealwithboundary中处理完毕
        for (SpMat::InnerIterator it(_A, k); it; ++it)
        {
            int row = it.row();
            if (row == k)
                continue;
            // if (level == 1)
            //     _Rhs[row] -= _A.coeffRef(k, row) * bnd_value;
            _A.coeffRef(k, row) = 0.0;
            _A.coeffRef(row, k) = 0.0;
        }
    }
}

int MultiGridMethod::switcher(int i, int j, int _nx, int _ny)
{
    if (j == 0)
    {
        if (i == 0)
            return 0;
        if (i == _nx - 1)
            return 2;
        return 1;
    }
    if (j == _ny - 1)
    {
        if (i == 0)
            return 6;
        if (i == _nx - 1)
            return 8;
        return 7;
    }
    if (i == 0)
        return 3;
    if (i == _nx - 1)
        return 5;
    return 4;
}

class Q1_MG : public MultiGridMethod
{
protected:
    typedef MultiGridMethod MG;
    typedef Eigen::Triplet<double> Tri;
    using MG::A;
    using MG::f;
    using MG::level;
    using MG::maxlevel;
    using MG::mesh;
    using MG::num;
    using MG::Solver_sparse;
    using MG::u;
    double weight[9][9] = {{1, 0.5, 0.5, 0.25}, {0.5, 1, 0.5, 0.25, 0.5, 0.25}, {0.5, 1, 0.25, 0.5}, {0.5, 0.25, 1, 0.5, 0.5, 0.25}, {0.25, 0.5, 0.25, 0.5, 1, 0.5, 0.25, 0.5, 0.25}, {0.25, 0.5, 0.5, 1, 0.25, 0.5}, {0.5, 0.25, 1, 0.5}, {0.25, 0.5, 0.25, 0.5, 1, 0.5}, {0.25, 0.5, 0.5, 1}};

public:
    Q1_MG(){};
    Q1_MG(SpMat &_A, RHS &_f, int _maxlevel, Mesh<2> *_mesh);
    ~Q1_MG(){};
    void V_Cycle(int _level);
    void reinit();
    vector<int> getStructure(int i, int j, int _nx, int _ny);
};

Q1_MG::Q1_MG(SpMat &_A, RHS &_f, int _maxlevel, Mesh<2> *_mesh)
{
    mesh = _mesh;
    level = 1;
    maxlevel = _maxlevel;
    A.push_back(_A);
    f.push_back(_f);
    u.push_back(Eigen::MatrixXd::Zero(f[0].size(), 1));
    A.resize(maxlevel);
    // u.resize(maxlevel);
    // f.resize(maxlevel);
    Solver_sparse.setTolerance(1e-12);
    reinit();
}

vector<int> Q1_MG::getStructure(int i, int j, int _nx, int _ny)
{

    int s = switcher(i, j, _nx, _ny);
    int idx = i + _nx * j;
    vector<int> structure;
    if (s != 0 && s != 1 && s != 2 && s != 3 && s != 6)
        structure.push_back(idx - _nx - 1);
    if (s != 0 && s != 1 && s != 2)
        structure.push_back(idx - _nx);
    if (s != 0 && s != 1 && s != 2 && s != 5 && s != 8)
        structure.push_back(idx - _nx + 1);
    if (s != 0 && s != 3 && s != 6)
        structure.push_back(idx - 1);
    structure.push_back(idx);
    if (s != 2 && s != 5 && s != 8)
        structure.push_back(idx + 1);
    if (s != 0 && s != 3 && s != 6 && s != 7 && s != 8)
        structure.push_back(idx + _nx - 1);
    if (s != 6 && s != 7 && s != 8)
        structure.push_back(idx + _nx);
    if (s != 2 && s != 5 && s != 6 && s != 7 && s != 8)
        structure.push_back(idx + _nx + 1);

    return structure;
}

void Q1_MG::V_Cycle(int _level)
{
    // vector<int> t = getStructure(9, 5, 5);
    // for (auto &it : t)
    //     cout << it << endl;
    int idx = 0;
    int pow2_level = POW2(_level);
    vector<int> indexofBoundaryDofs;
    for (int i = 0; i < mesh->getSegment()[0] / pow2_level + 1; i++)
        indexofBoundaryDofs.push_back(i);
    for (int j = 1; j < mesh->getSegment()[1] / pow2_level; j++)
    {
        indexofBoundaryDofs.push_back(j * (mesh->getSegment()[0] / pow2_level + 1));
        indexofBoundaryDofs.push_back((j + 1) * (mesh->getSegment()[0] / pow2_level + 1) - 1);
    }
    for (int i = 0; i < mesh->getSegment()[0] / pow2_level + 1; i++)
        indexofBoundaryDofs.push_back(i + (mesh->getSegment()[0] / pow2_level + 1) * mesh->getSegment()[1] / pow2_level);
    // for (int i = 0; i < indexofBoundaryDofs.size(); i++)
    //     cout << indexofBoundaryDofs[i] << endl;
    // cout << "--------------------------" << endl;
    level = _level;
    SpMat A_h = A[level - 1];
    RHS f_h = f[level - 1];
    DealWithBoundaryCondition(A_h, f_h);
    if (level == maxlevel)
    {
        Solver_sparse.compute(A_h);
        u[level - 1] = Solver_sparse.solve(f_h);
    }
    else
    {
        int n_h;
        if (level == 1)
            n_h = f[0].size();
        else
            n_h = A[level - 1].cols();
        int n_2h = ((sqrt(n_h) + 1) / 2) * ((sqrt(n_h) + 1) / 2);
        int nx = sqrt(n_h);
        int ny = sqrt(n_h);
        double *p;
        double *q[9];

        Smooth(A_h, u[level - 1], f_h, "GS", 2, mesh);
        RHS r_h = f_h - A_h * u[level - 1];

        // 计算r_2h
        // RHS r_2h = R[level - 1] * r_h;
        int row_2h = 0;
        RHS r_2h = MatrixXd::Zero(n_2h, 1);
        for (int j = 0; j < ny; j = j + 2)
        {
            for (int i = 0; i < nx; i = i + 2)
            {
                int s = switcher(i, j, nx, ny);
                p = &weight[s][0];
                for (auto &it : getStructure(i, j, nx, ny))
                    r_2h[row_2h] += r_h[it] * (*p++);
                row_2h++;
            }
        }

        // 令边界误差强制为0
        for (int i = 0; i < r_2h.size(); i++)
            if (i == indexofBoundaryDofs[idx])
            {
                r_2h[i] = 0.0;
                idx++;
            }
        // for (int i = 0; i < r_2h.size(); i++)
        //     cout << r_2h[i] << endl;
        // cout << "------------------" << endl;

        Solution e_2h = MatrixXd::Zero(r_2h.size(), 1);
        u.push_back(e_2h);
        f.push_back(r_2h);
        level++;
        V_Cycle(level);
        level--;

        // 计算e_2h
        // u[level - 1] = u[level - 1] + P[level - 1] * u[level];
        row_2h = 0;
        Solution e_h = MatrixXd::Zero(n_h, 1);
        for (int j = 0; j < ny; j = j + 2)
        {
            for (int i = 0; i < nx; i = i + 2)
            {
                int s = switcher(i, j, nx, ny);
                q[s] = &weight[s][0];
                for (auto &row_h : getStructure(i, j, nx, ny))
                    e_h[row_h] += u[level][row_2h] * (*q[s]++);
                row_2h++;
            }
        }
        u[level - 1] = u[level - 1] + e_h;
        Smooth(A_h, u[level - 1], f_h, "GS", 2, mesh);
    }
}

void Q1_MG::reinit()
{
    while (level < maxlevel)
    {
        int n_h;
        if (level == 1)
            n_h = f[0].size();
        else
            n_h = A[level - 1].cols();
        int n_2h = ((sqrt(n_h) + 1) / 2) * ((sqrt(n_h) + 1) / 2);
        int nx = sqrt(n_h);
        int ny = sqrt(n_h);
        double *p;
        double *q[9];
        // 计算A_2h
        // A[level] = R[level - 1] * A[level - 1] * P[level - 1];
        int row_2h = 0;
        A[level].resize(n_2h, n_2h);
        // A[level].setZero();
        clock_t start, end;
        start = clock();
        std::vector<Tri> TriList(n_2h * 9);
        std::vector<Tri>::iterator it = TriList.begin();
        for (int j = 0; j < ny; j += 2)
        {

            for (int i = 0; i < nx; i += 2)
            {
                clock_t start, end;

                // 计算A_2h[row_2h][getStructure(i/2,j/2,(nx+1)/2,(ny+1)/2)的每个元素]，这是为了保持A的稀疏性和提高效率
                vector<int> r;    // 存储A中每一列需要对应相乘的行号，每个i，j对应一个行号组
                vector<double> v; // 存储矩阵连乘的中间行向量，为下一个矩阵乘法做准备
                v.resize(n_h);
                for (auto &row_h : getStructure(i, j, nx, ny))
                    r.push_back(row_h);
                start = clock();
                for (int col_h = 0; col_h < n_h; col_h++)
                {
                    int s = switcher(i, j, nx, ny);
                    p = &weight[s][0];
                    for (auto &it : r)
                        v[col_h] += A[level - 1].coeff(it, col_h) * (*p++);
                }
                end = clock();
                double time = (double)(end - start) / CLOCKS_PER_SEC;
                cout << time << ", ";
                start = clock();
                for (auto &col_2h : getStructure(i / 2, j / 2, (nx + 1) / 2, (ny + 1) / 2))
                {
                    int x = col_2h % ((nx + 1) / 2);
                    int y = col_2h / ((nx + 1) / 2);
                    int s = switcher(2 * x, 2 * y, nx, ny);
                    p = &weight[s][0];
                    double sum = 0.0;
                    for (auto &col_h : getStructure(2 * x, 2 * y, nx, ny))
                        // A[level].coeffRef(row_2h, col_2h) += v[col_h] * (*p++);
                        sum += v[col_h] * (*p++);
                    *(it++) = Tri(row_2h, col_2h, sum);
                }
                end = clock();
                time = (double)(end - start) / CLOCKS_PER_SEC;
                cout << time << endl;
                row_2h++;
            }
        }
        end = clock();
        double time = (double)(end - start) / CLOCKS_PER_SEC;
        // cout << time << endl;
        A[level].setFromTriplets(TriList.begin(), TriList.end());
        A[level].makeCompressed();
        //  cout << A[level - 1] << endl;
        //  cout << A[level] << endl;
        level++;
    }
    level = 1;
}
#else
#endif
